package limiter;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 计数器限流
 * 固定时间窗口
 *
 * 缺点：临界点附近可能会有大量并发被允许访问，对系统造成冲击
 */
public class CountLimiterDemo {
    public static void main(String[] args) throws InterruptedException {
        final CountLimit limit = new CountLimit(10, 1, TimeUnit.SECONDS);
        ExecutorService executorService = Executors.newCachedThreadPool();
        Random random = new Random();
        for (int i = 0; i < 100; i++) {
            final int f = i;
            executorService.execute(() -> {
                if(limit.tryAcquire())
                    System.out.println(new SimpleDateFormat("mm:ss.SSS").format(new Date()));
//                else
//                    System.out.println(Thread.currentThread() + " 被限流了");
            });
            Thread.sleep(random.nextInt(60));
        }
        executorService.shutdown();
    }

    private static class CountLimit {
        /**
         * 一定时间内允许请求的数量
         */
        private final long limit;
        /**
         * 间隔时间，毫秒
         */
        private final long interval;

        private AtomicLong startTime;
        private AtomicLong count = new AtomicLong(0);

        public CountLimit(long limit, long interval, TimeUnit timeUnit) {
            assert limit > 0;
            assert interval > 0;
            this.limit = limit;
            this.interval = TimeUnit.MILLISECONDS.convert(interval, timeUnit);
            startTime = new AtomicLong(System.currentTimeMillis());
        }

        public boolean tryAcquire() {
            for(;;) {
                long st = startTime.get();
                if (System.currentTimeMillis() - st < interval) { // 在统计周期内
                    long ct;
                    do {
                        ct = count.get();
                        if (ct >= limit) // 已经到达了限制，则限流
                            return false;
                    } while (!count.compareAndSet(ct, ct + 1));
                    return true;
                } else {
                    if (startTime.compareAndSet(st, System.currentTimeMillis())) {// 只有一个线程能修改成功
                        count.set(1); // 清空计数器
                        return true;
                    }
                }
            }
        }
    }
}
